# coding: utf-8

from __future__ import absolute_import, division, print_function, unicode_literals

from thrift2pyi.convert import Thrift2pyi


def get_type(func):
    def wrapped(self, args):
        _type = func(self, args)
        if isinstance(args, tuple):
            thrift_type, nest = args
            try:
                if nest.__module__ == self.thrift.__name__:
                    return nest.__name__
                else:
                    module = nest.__module__
                    self._imports[self._module2package[module]].add(("%s_thrift" % module, module))
                    return "%s.%s" % (module, nest.__name__)
            except AttributeError:
                return _type
        else:
            return _type

    return wrapped


Thrift2pyi._get_type = get_type(getattr(Thrift2pyi, "_get_type"))
